/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2012-2025. All rights reserved.
 * Create Date : 2025
 */

#ifndef TRANSPORT_HETEROG_PD_ROCE_PUB_H
#define TRANSPORT_HETEROG_PD_ROCE_PUB_H

#include "transport_heterog_roce_pub.h"

namespace hccl {
class TransportHeterogRawRoce : public TransportHeterogRoce {
public:
    explicit TransportHeterogRawRoce(const std::string &transTag, HcclIpAddress &selfIp,
        HcclIpAddress &peerIp, u32 peerPort, u32 selfPort, const TransportResourceInfo &transportResourceInfo);
    ~TransportHeterogRawRoce() override;

    HcclResult Init() override;
    HcclResult Init(socket_info_t &socketInfo, RdmaHandle rdmaHandle, MrHandle mrHandle) override;
    HcclResult ImrecvScatter(void *buf[], int count[], int bufCount, HcclDataType datatype, HcclMessageInfo &msg,
        HcclRequestInfo *&request) override;

protected:
    HcclResult CreateCqAndQp() override;
    HcclResult DestroyCqAndQp() override;
    HcclResult PreQpConnect() override;

private:
    HcclResult EnterStateProcess(ConnState nextState) override;
    HcclResult LoopStateProcess() override;
    HcclResult PrepareModifyInfo(struct qp_attr &qpAttr, struct typical_qp &typicalQpInfo);
    HcclResult GetQpAttr(QpHandle &qpHandle, struct qp_attr *attr, bool &completed);
    HcclResult TypicalQpModify(QpHandle &qpHandle, struct typical_qp* localQpInfo,
        struct typical_qp* remoteQpInfo, bool &completed);

    struct ibv_send_wr dataReadWrScatter_;      // scatter数据读取的wr模板
    struct ibv_send_wr dataAckWrScatter_;       // scatterACK发送的wr模板

    struct qp_attr localTagQpAttr_{0};
    struct qp_attr localDataQpAttr_{0};
    struct qp_attr remoteTagQpAttr_{0};
    struct qp_attr remoteDataQpAttr_{0};
    struct typical_qp localTagModifyInfo_{0};
    struct typical_qp localDataModifyInfo_{0};
    struct typical_qp remoteTagModifyInfo_{0};
    struct typical_qp remoteDataModifyInfo_{0};
};
}
#endif